from torch import nn
from torchinfo import summary

class Classifier(nn.Module):
    def __init__(self):
        super(Classifier, self).__init__()
        # torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        # torch.nn.MaxPool2d(kernel_size, stride, padding)
        # input 維度 [3, 128, 128]
        self.cnn = nn.Sequential(
            nn.Conv2d(3, 64, 3, 1, 1),  # [64, 128, 128]
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2, 2, 0),      # [64, 64, 64]

            nn.Conv2d(64, 128, 3, 1, 1), # [128, 64, 64]
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2, 2, 0),      # [128, 32, 32]

            nn.Conv2d(128, 256, 3, 1, 1), # [256, 32, 32]
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(2, 2, 0),      # [256, 16, 16]

            nn.Conv2d(256, 512, 3, 1, 1), # [512, 16, 16]
            nn.BatchNorm2d(512),
            nn.ReLU(),
            nn.MaxPool2d(2, 2, 0),       # [512, 8, 8]

            nn.Conv2d(512, 512, 3, 1, 1), # [512, 8, 8]
            nn.BatchNorm2d(512),
            nn.ReLU(),
            nn.MaxPool2d(2, 2, 0),       # [512, 4, 4]
        )
        self.fc = nn.Sequential(
            nn.Linear(512*4*4, 1024),
            nn.ReLU(),
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, 11)
        )

    def forward(self, x):
        out = self.cnn(x)
        out = out.view(out.size()[0], -1)
        return self.fc(out)

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride = 1, downsample=None):
        super(ResidualBlock, self).__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, stride, 1),
            nn.BatchNorm2d(out_channels),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
        )
        self.relu = nn.ReLU()
        self.downsample = downsample  # 用于调整维度以匹配残差连接

    def forward(self, x):
        identity = x  # 残差连接的输入
        out = self.cnn(x)

        # 如果需要，调整输入以匹配输出维度
        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity  # 残差连接
        out = self.relu(out)
        return out


class ResidualNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=11):
        super(ResidualNet, self).__init__()
        self.pre_layer = nn.Sequential(
            nn.Conv2d(in_channels, 64, 7, 2, 3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(3, 2, 1)
        )
        self.layer1 = self._make_layer(64, 64, 1)
        self.layer2 = self._make_layer(64, 128, 2)
        self.layer2_1 = self._make_layer(128, 128, 1)
        self.layer3 = self._make_layer(128, 256, 2)
        self.layer3_1 = self._make_layer(256, 256, 1)
        self.layer4 = self._make_layer(256, 512, 2)
        self.layer4_1 = self._make_layer(512, 512, 1)
        self.avg_pool = nn.AvgPool2d(3, 2, 1)
        self.fc = nn.Sequential(
            nn.Linear(512*4*4, 2048),
            nn.ReLU(),
            nn.Linear(2048, 1024),
            nn.ReLU(),
            nn.Linear(1024, out_channels)
        )

    def _make_layer(self, in_channels, out_channels, stride=1):
        downsample = None
        # 如果输入和输出通道数不匹配或者步长不为1，需要下采样
        if in_channels != out_channels:
            downsample = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, stride),
                nn.BatchNorm2d(out_channels)
            ) 
        return ResidualBlock(in_channels, out_channels, stride, downsample)



    def forward(self, x):
        x = self.pre_layer(x)   # [64, 64, 64]
        # stage 1
        # for _ in range(3):
        #     x = self.layer1(x)  # [64, 64, 64]
        x = self.layer1(self.layer1(x))      # [64, 64, 64]

        # stage 2
        x = self.layer2(x)      # [128, 32, 32]
        # for _ in range(3):
        #     x = self.layer2_1(x)  # [128, 32, 32]
        x = self.layer2_1(x)    # [128, 32, 32]

        # stage 3
        x = self.layer3(x)      # [256, 16, 16]
        # for _ in range(6):
        #     x = self.layer3_1(x)  # [256, 16, 16]
        x = self.layer3_1(x)    # [256, 16, 16]

        
        # stage 4
        x = self.layer4(x)      # [512, 8, 8]
        # for _ in range(2):
        #     x = self.layer4_1(x)  # [512, 8, 8]
        x = self.layer4_1(x)    # [512, 8, 8]

        x = self.avg_pool(x)     # [512, 4, 4]
        x = x.view(x.size()[0], -1)
        return self.fc(x)


if __name__ == "__main__":
    model_res = ResidualNet()
    summary(model_res, input_size=(64, 3, 256, 256))

    # model_cnn = Classifier()
    # summary(model_cnn, input_size=(64, 3, 128, 128))
    
